package com.xyzwps.example.stream;

import java.util.Objects;
import java.util.function.Function;

public class StreamFlatMapOperator<T, R> implements Operator<T, R> {

    private final Function<T, Stream<R>> flatMapFn;


    public StreamFlatMapOperator(Function<T, Stream<R>> flatMapFn) {
        this.flatMapFn = Objects.requireNonNull(flatMapFn);
    }

    @Override
    public Stream<R> handle(Stream<T> origin) {
        var self = this;
        return new Stream<>() {
            private Stream<R> subStream;

            @Override
            public Parcel<R> next() {
                if (this.subStream == null) {
                    switch (origin.next()) {
                        case Parcel.Done<T> ignored -> {
                            return Parcel.done();
                        }
                        case Parcel.Failure<T> failure -> {
                            return Parcel.failed(failure.cause());
                        }
                        case Parcel.Delivery<T> delivery -> this.subStream = self.flatMapFn.apply(delivery.value());
                    }
                }

                switch (subStream.next()) {
                    case Parcel.Done<R> ignored -> {
                        this.subStream = null;
                        return this.next(); // TODO: 递归了。。。
                    }
                    case Parcel.Failure<R> failure -> {
                        return Parcel.failed(failure.cause());
                    }
                    case Parcel.Delivery<R> delivery -> {
                        return Parcel.deliver(delivery.value());
                    }
                }
                throw new RuntimeException("TODO: impossible");
            }
        };
    }
}
